from nju.comp.linearGenerator import Generator
from nju.comp.linearDiscriminator import Discriminator
from nju.utils.LossComputeWithModule import LossComputeWithModule
import torch
import torch.nn as nn
import numpy as np


class LinearGAN(nn.Module):
    def __init__(self):
        super(LinearGAN, self).__init__()
        self.gen = Generator()
        self.dis = Discriminator()
        self.gen_optim = torch.optim.SGD(self.gen.parameters(), lr=0.1)
        self.dis_optim = torch.optim.SGD(self.dis.parameters(), lr=0.1)
        self.criterion = LossComputeWithModule()


